-- 令牌桶限流: 支持预消费, 初始桶是满的
-- KEYS[1]  string  限流的key

-- ARGV[1]  int     桶最大容量
-- ARGV[2]  int     每次添加令牌数
-- ARGV[3]  int     令牌添加间隔(秒)
-- ARGV[4]  int     想获取令牌的数量
-- ARGV[5]  int     是否vip 0-否  1-时
local bucket_capacity = tonumber(ARGV[1])
local add_token = tonumber(ARGV[2])
local add_interval = tonumber(ARGV[3])
local expect_count = tonumber(ARGV[4])
local is_vip = tonumber(ARGV[5])
-- 当前时间（s）
local now = redis.call('TIME')[1]

-- 保存上一次更新桶的时间的key
local LAST_TIME_KEY = "{"..KEYS[1].."}_time";
-- 获取当前桶中令牌数
local token_cnt = redis.call("get", KEYS[1])
-- 最终返回值，-1代表未获取到令牌
local returnCount = -1


if token_cnt then   -- 令牌桶存在
    -- 上一次更新桶的时间（s）
    local last_time = redis.call('get', LAST_TIME_KEY)
    -- 恢复倍数
    local multiple = math.floor((now - last_time) / add_interval)
    -- 这段时间可以恢复令牌数
    local recovery_cnt = multiple * add_token
    -- 计算桶当前令牌数，确保不超过桶容量
    token_cnt = math.min(bucket_capacity, token_cnt + recovery_cnt)
else
	token_cnt = bucket_capacity
end

-- 非vip，桶内无令牌，返回-1
if is_vip==0 and token_cnt <= 0 then
    return returnCount;
-- 非vip，桶内有令牌，返回 桶内令牌和期望值得最小值
elseif is_vip==0 and token_cnt > 0 then
    returnCount = math.min(token_cnt,expect_count)
    token_cnt = token_cnt - returnCount
-- vip 允许预消费
elseif is_vip==1 then
    returnCount = expect_count
    token_cnt = token_cnt - expect_count
end

-- 桶完全恢复需要的最大时长（s），由于允许预消费，所有桶内可能为负值
local reset_time = math.ceil( (bucket_capacity - token_cnt) / add_token) * add_interval

-- 设置过期时间避免key一直存在
redis.call('set', KEYS[1], token_cnt, 'EX', reset_time)
redis.call('set', LAST_TIME_KEY, now, 'EX', reset_time)
return returnCount